import torch

import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import SAGEConv, to_hetero
dataset = OGB_MAG(root='/home/Dyf/code/dataset/data', preprocess='metapath2vec',transform=T.ToUndirected())
data = dataset[0]
from torch_geometric.nn import GATConv, Linear, to_hetero

class GAT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GATConv((-1, -1), hidden_channels, add_self_loops=False)
        self.lin1 = Linear(-1, hidden_channels)
        self.conv2 = GATConv((-1, -1), out_channels, add_self_loops=False)
        self.lin2 = Linear(-1, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index) + self.lin1(x)
        x = x.relu()
        x = self.conv2(x, edge_index) + self.lin2(x)
        return x

model = GAT(hidden_channels=64, out_channels=dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')
print(model)
# print(model.print_readable())

"""
GraphModule(
  (conv1): ModuleDict(
    (author__affiliated_with__institution): GATConv((-1, -1), 64, heads=1)
    (author__writes__paper): GATConv((-1, -1), 64, heads=1)
    (paper__cites__paper): GATConv((-1, -1), 64, heads=1)
    (paper__has_topic__field_of_study): GATConv((-1, -1), 64, heads=1)
    (institution__rev_affiliated_with__author): GATConv((-1, -1), 64, heads=1)
    (paper__rev_writes__author): GATConv((-1, -1), 64, heads=1)
    (field_of_study__rev_has_topic__paper): GATConv((-1, -1), 64, heads=1)
  )
  (lin1): ModuleDict(
    (paper): Linear(-1, 64, bias=True)
    (author): Linear(-1, 64, bias=True)
    (institution): Linear(-1, 64, bias=True)
    (field_of_study): Linear(-1, 64, bias=True)
  )
  (conv2): ModuleDict(
    (author__affiliated_with__institution): GATConv((-1, -1), 349, heads=1)
    (author__writes__paper): GATConv((-1, -1), 349, heads=1)
    (paper__cites__paper): GATConv((-1, -1), 349, heads=1)
    (paper__has_topic__field_of_study): GATConv((-1, -1), 349, heads=1)
    (institution__rev_affiliated_with__author): GATConv((-1, -1), 349, heads=1)
    (paper__rev_writes__author): GATConv((-1, -1), 349, heads=1)
    (field_of_study__rev_has_topic__paper): GATConv((-1, -1), 349, heads=1)
  )
  (lin2): ModuleDict(
    (paper): Linear(-1, 349, bias=True)
    (author): Linear(-1, 349, bias=True)
    (institution): Linear(-1, 349, bias=True)
    (field_of_study): Linear(-1, 349, bias=True)
  )
)
def forward(self, x, edge_index):
    x_dict = torch_geometric_nn_to_hetero_transformer_get_dict(x);  x = None
    x__paper = x_dict.get('paper', None)
    x__author = x_dict.get('author', None)
    x__institution = x_dict.get('institution', None)
    x__field_of_study = x_dict.get('field_of_study', None);  x_dict = None
    
    # 以下的这一行计算了所有的索引张量并组成一个字典形式用于进行不同边的聚合运算。
    edge_index_dict = torch_geometric_nn_to_hetero_transformer_get_dict(edge_index);  edge_index = None
    edge_index__author__affiliated_with__institution = edge_index_dict.get(('author', 'affiliated_with', 'institution'), None)
    edge_index__author__writes__paper = edge_index_dict.get(('author', 'writes', 'paper'), None)
    edge_index__paper__cites__paper = edge_index_dict.get(('paper', 'cites', 'paper'), None)
    edge_index__paper__has_topic__field_of_study = edge_index_dict.get(('paper', 'has_topic', 'field_of_study'), None)
    edge_index__institution__rev_affiliated_with__author = edge_index_dict.get(('institution', 'rev_affiliated_with', 'author'), None)
    edge_index__paper__rev_writes__author = edge_index_dict.get(('paper', 'rev_writes', 'author'), None)
    edge_index__field_of_study__rev_has_topic__paper = edge_index_dict.get(('field_of_study', 'rev_has_topic', 'paper'), None);  edge_index_dict = None
    
    conv1__institution = self.conv1.author__affiliated_with__institution((x__author, x__institution), edge_index__author__affiliated_with__institution)
    conv1__paper1 = self.conv1.author__writes__paper((x__author, x__paper), edge_index__author__writes__paper)
    conv1__paper2 = self.conv1.paper__cites__paper(x__paper, edge_index__paper__cites__paper)
    conv1__field_of_study = self.conv1.paper__has_topic__field_of_study((x__paper, x__field_of_study), edge_index__paper__has_topic__field_of_study)
    conv1__author1 = self.conv1.institution__rev_affiliated_with__author((x__institution, x__author), edge_index__institution__rev_affiliated_with__author)
    conv1__author2 = self.conv1.paper__rev_writes__author((x__paper, x__author), edge_index__paper__rev_writes__author)
    conv1__paper3 = self.conv1.field_of_study__rev_has_topic__paper((x__field_of_study, x__paper), edge_index__field_of_study__rev_has_topic__paper)
    conv1__paper_1 = torch.add(conv1__paper1, conv1__paper2);  conv1__paper1 = conv1__paper2 = None
    conv1__paper = torch.add(conv1__paper3, conv1__paper_1);  conv1__paper3 = conv1__paper_1 = None
    conv1__author = torch.add(conv1__author1, conv1__author2);  conv1__author1 = conv1__author2 = None
    lin1__paper = self.lin1.paper(x__paper);  x__paper = None
    lin1__author = self.lin1.author(x__author);  x__author = None
    lin1__institution = self.lin1.institution(x__institution);  x__institution = None
    lin1__field_of_study = self.lin1.field_of_study(x__field_of_study);  x__field_of_study = None
    add__paper = conv1__paper + lin1__paper;  conv1__paper = lin1__paper = None
    add__author = conv1__author + lin1__author;  conv1__author = lin1__author = None
    add__institution = conv1__institution + lin1__institution;  conv1__institution = lin1__institution = None
    add__field_of_study = conv1__field_of_study + lin1__field_of_study;  conv1__field_of_study = lin1__field_of_study = None
    relu__paper = add__paper.relu();  add__paper = None
    relu__author = add__author.relu();  add__author = None
    relu__institution = add__institution.relu();  add__institution = None
    relu__field_of_study = add__field_of_study.relu();  add__field_of_study = None
    conv2__institution = self.conv2.author__affiliated_with__institution((relu__author, relu__institution), edge_index__author__affiliated_with__institution);  edge_index__author__affiliated_with__institution = None
    conv2__paper1 = self.conv2.author__writes__paper((relu__author, relu__paper), edge_index__author__writes__paper);  edge_index__author__writes__paper = None
    conv2__paper2 = self.conv2.paper__cites__paper(relu__paper, edge_index__paper__cites__paper);  edge_index__paper__cites__paper = None
    conv2__field_of_study = self.conv2.paper__has_topic__field_of_study((relu__paper, relu__field_of_study), edge_index__paper__has_topic__field_of_study);  edge_index__paper__has_topic__field_of_study = None
    conv2__author1 = self.conv2.institution__rev_affiliated_with__author((relu__institution, relu__author), edge_index__institution__rev_affiliated_with__author);  edge_index__institution__rev_affiliated_with__author = None
    conv2__author2 = self.conv2.paper__rev_writes__author((relu__paper, relu__author), edge_index__paper__rev_writes__author);  edge_index__paper__rev_writes__author = None
    conv2__paper3 = self.conv2.field_of_study__rev_has_topic__paper((relu__field_of_study, relu__paper), edge_index__field_of_study__rev_has_topic__paper);  edge_index__field_of_study__rev_has_topic__paper = None
    conv2__paper_1 = torch.add(conv2__paper1, conv2__paper2);  conv2__paper1 = conv2__paper2 = None
    conv2__paper = torch.add(conv2__paper3, conv2__paper_1);  conv2__paper3 = conv2__paper_1 = None
    conv2__author = torch.add(conv2__author1, conv2__author2);  conv2__author1 = conv2__author2 = None
    lin2__paper = self.lin2.paper(relu__paper);  relu__paper = None
    lin2__author = self.lin2.author(relu__author);  relu__author = None
    lin2__institution = self.lin2.institution(relu__institution);  relu__institution = None
    lin2__field_of_study = self.lin2.field_of_study(relu__field_of_study);  relu__field_of_study = None
    add_1__paper = conv2__paper + lin2__paper;  conv2__paper = lin2__paper = None
    add_1__author = conv2__author + lin2__author;  conv2__author = lin2__author = None
    add_1__institution = conv2__institution + lin2__institution;  conv2__institution = lin2__institution = None
    add_1__field_of_study = conv2__field_of_study + lin2__field_of_study;  conv2__field_of_study = lin2__field_of_study = None
    return {'paper': add_1__paper, 'author': add_1__author, 'institution': add_1__institution, 'field_of_study': add_1__field_of_study}
    

1、基于边进行聚合运算。
2、计算的结果sum汇总到指向的目的节点。
3、计算的顺序相当重要。必须将其中一层的信息都计算之后，才能进行特征融合，然后才可以进行下一层计算。避免混乱。
4、计算过程中是基于目的节点的汇总，每一个卷积层第二个参数，是融合目的地。
"""